package cn.foolishbird.crow.core.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

/**
 * @author foolish-bird
 */
public final class CollectionTools {

    private CollectionTools() {

    }

    /**
     * 两个集合求并集
     *
     * @param a
     * @param b
     * @param apply
     * @param <O>
     * @param <R>
     * @return
     */
    public static <O, R> Collection<O> union(Collection<? extends O> a, Collection<? extends O> b,
                                             Function<? extends O, ? extends R> apply) {
        CollectionTools.SetOperationCardinalityHelper<O, R> helper = new CollectionTools.SetOperationCardinalityHelper(
            a, b, apply);
        Iterator<? extends O> iterator = helper.iterator();

        while (iterator.hasNext()) {
            O obj = iterator.next();
            helper.setCardinality(obj, helper.max(obj));
        }

        return helper.list();
    }

    /**
     * 集合求交集
     *
     * @param a
     * @param b
     * @param apply
     * @param <O>
     * @param <R>
     * @return
     */
    public static <O, R> Collection<O> intersection(Collection<? extends O> a, Collection<? extends O> b,
                                                    Function<? extends O, ? extends R> apply) {
        CollectionTools.SetOperationCardinalityHelper<O, R> helper = new CollectionTools.SetOperationCardinalityHelper(
            a, b, apply);
        Iterator<O> iterator = helper.iterator();

        while (iterator.hasNext()) {
            O obj = iterator.next();
            helper.setCardinality(obj, helper.min(obj));
        }

        return helper.list();
    }

    /**
     * 集合A不集合B不交的
     *
     * @param a
     * @param b
     * @param apply
     * @param <O>
     * @param <R>
     * @return
     */
    public static <O, R> Collection<O> disjunction(Collection<? extends O> a, Collection<? extends O> b,
                                                   Function<? extends O, ? extends R> apply) {
        CollectionTools.SetOperationCardinalityHelper<O, R> helper = new CollectionTools.SetOperationCardinalityHelper(
            a, b, apply);
        Iterator<O> iterator = helper.iterator();

        while (iterator.hasNext()) {
            O obj = iterator.next();
            helper.setCardinality(obj, helper.max(obj) - helper.min(obj));
        }

        return helper.list();
    }

    /**
     * 集合A和集合B，求差集
     *
     * @param a
     * @param b
     * @param f
     * @param <O>
     * @param <R>
     * @return
     */
    public static <O, R> Collection<O> subtract(Collection<? extends O> a, Collection<? extends O> b, Function<O, R> f) {
        Map<R, Integer> bag = new HashMap<>(16);
        O element = null;
        Iterator<? extends O> iterator = b.iterator();
        while (iterator.hasNext()) {
            element = iterator.next();
            bag.put(f.apply(element), 1);
        }

        List<O> subtract = new ArrayList();
        iterator = a.iterator();
        while (iterator.hasNext()) {
            element = iterator.next();
            if (!bag.containsKey(f.apply(element))) {
                subtract.add(element);
            }
        }

        return subtract;
    }

    /**
     * 辅助集合操作
     *
     * @param <O>
     * @param <R>
     */
    private static class SetOperationCardinalityHelper<O, R> extends CollectionTools.CardinalityHelper<O, R>
        implements Iterable<O> {
        private final   List<O> elements;
        private final Set<R> keys;
        private final List<O> newList;

        public SetOperationCardinalityHelper(Collection<? extends O> a, Collection<? extends O> b,
                                             Function<O, R> apply) {
            super(a, b, apply);
            // initialize container
            int size = a.size() + b.size();
            this.elements = Lists.newArrayListWithCapacity(size);
            this.keys = Sets.newHashSetWithExpectedSize(size);

            this.elements.addAll(a.stream().filter(e -> keys.add(apply.apply(e))).collect(Collectors.toList()));
            this.elements.addAll(b.stream().filter(e -> keys.add(apply.apply(e))).collect(Collectors.toList()));
            this.newList = new ArrayList(this.elements.size());
        }

        public Iterator<O> iterator() {
            return this.elements.iterator();
        }

        public void setCardinality(O obj, int count) {
            for (int i = 0; i < count; ++i) {
                this.newList.add(obj);
            }

        }

        public Collection<O> list() {
            return this.newList;
        }
    }

    /**
     * 统计、查询集合元素数量
     *
     * @param <T>
     * @param <R>
     */
    private static class CardinalityHelper<T, R> {

        private final Map<R, Integer> cardinalityA;

        private final Map<R, Integer> cardinalityB;

        private final Function<T, R> apply;

        public CardinalityHelper(Collection<? extends T> a, Collection<? extends T> b,
                                 Function<T, R> apply) {
            this.cardinalityA = cardinalityMap(a, apply);
            this.cardinalityB = cardinalityMap(b, apply);
            this.apply = apply;
        }

        public final int max(T obj) {
            return Math.max(this.freqA(obj), this.freqB(obj));
        }

        public final int min(T obj) {
            return Math.min(this.freqA(obj), this.freqB(obj));
        }

        public int freqA(T obj) {
            return this.getFreq(obj, this.cardinalityA);
        }

        public int freqB(T obj) {
            return this.getFreq(obj, this.cardinalityB);
        }

        private final int getFreq(T obj, Map<R, Integer> freqMap) {
            Integer count = freqMap.get(this.apply.apply(obj));
            return count != null ? count : 0;
        }

    }

    /**
     * 统计集合元素数量
     *
     * @param coll
     * @param apply
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> Map<R, Integer> cardinalityMap(Collection<? extends T> coll,
                                                        Function<T, R> apply) {

        Map<R, Integer> count = Maps.newHashMapWithExpectedSize(coll.size());

        Iterator<? extends T> iterator = coll.iterator();
        while (iterator.hasNext()) {
            T obj = iterator.next();
            Integer c = count.get(obj);
            R key = apply.apply(obj);
            if (null == c) {
                count.put(key, 1);
            } else {
                count.put(key, c + 1);
            }
        }
        return count;
    }

}
